;;;; -*- Mode: Lisp; -*-

;;; ------------------------------------------------------
;;; Linear Algebra using matrices as defned in matrix.lisp
;;;
;;;


(in-package :lib.gl)

;;; General use (more so I understand the meaning of the notation in [Murty 97]
;;; than anything else)
(defun zero-vector-p (v)
  "T if the absolute value of all elements of <v> are less than +epsilon+."
  (iter (for x in-vector v)
	(always (< (abs x) +epsilon+))))

(defun positive-vector-p (v)
  "T if all elements of <v> are greater than 0.0."
  (iter (for x in-vector v)
	(always (plusp x))))

(defun semipositive-vector-p (v)
  "T if v is not the zero vector and all elements are greater than or equal to 0."
  (and (not (zero-vector-p v))
       (iter (for x in-vector v)
	     (always (or (plusp x) (< (abs x) +epsilon+))))))

(defun nonnegative-vector-p (v)
  "T if v is the zero vector, is positive, or is semipositive."
  (iter (for x in-vector v)
	(always (or (plusp x) (< (abs x) +epsilon+)))))

(defun diagonal-p (m)
  (let ((r (matrix-rows m))
	(c (matrix-cols m)))
    (iter (for i from 0 below r)
	  (always (iter (for j from 0 below c)
			(always (if (= i j)
				    (not (zerop (mref m i j)))
				    (zerop (mref m i j)))))))))

;;; Operations for gaussian elimination

(defun gaussian-row-exchange-slow (matrix ri rj)
  "In <matrix>, exchange row <ri> with row <rj>."
  (let* ((op (identity-matrix (matrix-rows matrix)
			      :element-type (element-type matrix)))
	 (coli (get-column op ri)) ;; We can do this because I = transpose(I)
	 (colj (get-column op rj)))
    (set-column op rj coli)
    (set-column op ri colj)
    (m* op matrix)))

(defun gaussian-row-exchange (matrix ri rj)
  "In <matrix>, exchange row <ri> with row <rj>."
  (let* ((row-i (get-row matrix ri))
	 (row-j (get-row matrix rj)))
    (set-row matrix rj row-i)
    (set-row matrix ri row-j)
    matrix))

(defun gaussian-row-scale-slow (matrix row c)
  "Multiply/scale a row, indicated by integer argument <row>, in the
given matrix by constant <c>."
  (declare (fixnum row))
  (let* ((op (identity-matrix (matrix-rows matrix)
			      :element-type (element-type matrix))))
    (setf (mref op row row) (coerce c (element-type matrix)))
    (m* op matrix)))

(defun gaussian-row-scale (matrix row c)
  "Multiply/scale a row, indicated by integer argument <row>, in the
given matrix by constant <c>."
  (iter (for j from 0 below (matrix-cols matrix))
	(setf (mref matrix row j) (* c (mref matrix row j))))
  matrix)

(defun gaussian-row-scale-sum-slow (matrix ri c rj)
  "Multiply row <ri> of matrix matrix by constant <c> and sum it with
row <rj>, replacing row <rj> with the sum."
  (let* ((op (identity-matrix (matrix-rows matrix) :element-type (element-type matrix))))
    (setf (mref op rj ri) (coerce c (element-type matrix)))
    (m* op matrix)))

(defun gaussian-row-scale-sum (matrix ri c rj)
  (let* ((row-i (get-row matrix ri))
	 (row-j (get-row matrix rj)))
    (iter (for j from 0 below (matrix-cols matrix))
	  (for xi in-vector row-i)
	  (for xj in-vector row-j)
	  (setf (mref matrix rj j) (+ (* c xi) xj)))
    matrix))
	 

(defun gaussian-pivot (matrix r s)
  "Perform a gaussian elimination on matrix <matrix> from the pivot
indicated by row <r> and column <s>."
  (gaussian-row-scale matrix r (/ 1 (mref matrix r s)))
  (let ((nrows (matrix-rows matrix)))
    (iter (for ri from 0 below nrows)
	  (unless (= ri r)
	    (gaussian-row-scale-sum matrix r (- (mref matrix ri s)) ri)))
    matrix))

(defun gaussian-forward-elimination (matrix)
  "Perform a gaussian forward elimination on <matrix>, destructively
modifying <matrix> in the process."
  (let ((nrows (matrix-rows matrix)))
    (iter (for ci from 0 below nrows)
	  (let ((pivot (mref matrix ci ci)))
	    (iter (for ri from (1+ ci) below nrows)
		  (gaussian-row-scale-sum matrix ci (- (/ (mref matrix ri ci) pivot)) ri))))
    matrix))
	  
(defun gaussian-back-substitution (matrix)
  "Back substitution of matrix <matrix> (presumably after doing a gaussian forward elimination),
to determine the solutions to the linear system described by
<matrix>."
  (let ((nrows (matrix-rows matrix)))
    (iter (for ci from (1- nrows) downto 1)
	  (let ((pivot (mref matrix ci ci)))
	    (setf matrix (gaussian-row-scale matrix ci (/ 1 pivot)))
	    (iter (for ri from (1- ci) downto 0)
		  (setf matrix (gaussian-row-scale-sum matrix ci (- (mref matrix ri ci)) ri)))))
    (setf matrix (gaussian-row-scale matrix 0 (/ 1 (mref matrix 0 0))))
    matrix))

;;; Based on Fortran code at http://www.engr.colostate.edu/~thompson/hPage/CourseMat/Tutorials/ComMethods/doolittle.pdf

(defun doolittle-lu-decomposition (a)
  "Find the LU decomposition of matrix <a> using the doolittle algorithm."
  (let* ((n (matrix-rows a))
	 (l (identity-matrix n :element-type (element-type a)))
	 (u (identity-matrix n :element-type (element-type a))))
    (dotimes (i n (values l u))
      (iter (for j from i below n)
	    (let ((a_ij (mref a i j)))
	      (setf (mref u i j) a_ij)
	      (unless (= i 0)
		(iter (for k from 0 to (1- i))
		      (setf (mref u i j) (- (mref u i j) (* (mref l i k) (mref u k j))))))))
      (iter (for j from (1+ i) below n)
	    (let ((a_ji (mref a j i)))
	      (setf (mref l j i) a_ji)
	      (unless (= i 0)
		(iter (for k from 0 to (1- i))
		      (setf (mref l j i) (- (mref l j i) (* (mref l j k) (mref u k i))))))
	      (setf (mref l j i) (/ (mref l j i) (mref u i i))))))))

    
;;;; ---- Orthogonal Matrices and Gram-Schmidt Orthogonalization ---

(defun orthogonalize-matrix (a)
  "Using Gram-Schmidt, produce an orthogonal matrix from matrix <a>."
  (let* ((cols (matrix-cols a))
	 (rows (matrix-rows a))
	 (q (make-matrix rows cols))
	 (a0 (get-column a 0))
	 (q0 (m/ a0 (magnitude a0))))
    (set-column q 0 q0)
    (iter (for j from 1 below cols)
	  (finally (return q))
	  (let ((aj (get-column a j))
		(psum (zero-vector rows)))
	    (iter (for n from 0 below j)
		  (let ((qn (get-column q n)))
		    (setf psum (m+ psum (m* (dot-product qn aj) qn)))))
	    (let ((qj (m- aj psum)))
	      (set-column q j (m/ qj (magnitude qj))))))))
		  
#|(defun fast-qr-factorization (a)
  "A faster algorithm for qr factorization taking advantage of the fact that
a is tridiagonal (A is tridiagonal, right? This algorithm won't check for you)."
  (let* ((nc (matrix-cols a))
	 (nr (matrix-rows a))
	 (q (make-matrix nr nc))
	 (a00 (mref a 0 0))
	 (a01 (mref a 1 0))
	 (m (/ 1 (sqrt (+ (* a00 a00) (a01 a01))))))
    (flet ((subrowsum (i end)
	     (iter (for j from 0 below end)
		   (summing (mref q i j))))
	   (subcolmag (j end)
	     (sqrt (iter (for i from 0 below end)
			 (with qij = (mref q i j))
			 (summing (* qij qij))))))
      (setf (mref q 0 0) (* a00 m))
      (setf (mref q 1 0) (* a01 m))
      (iter (for j from 1 below nc)
	    (finally (return q))
	    (iter (for i from 0 to j)
|#

(defun qr-factorization (a &optional (tridiagonal-p nil))
  (cond (tridiagonal-p (fast-qr-factorization))
	(t (let* ((q (orthogonalize-matrix a))
		  (cols (matrix-cols a))
		  (rows (matrix-rows a))
		  (r (make-matrix rows cols)))
	     (dotimes (i rows (values q r))
	       (dotimes (j cols)
		 (setf (mref r i j)
		       (dot-product (get-column q i)
				    (get-column a j)))))))))
	
;;;; ---- Properties of Matrices ----
(defun positive-definite-p (m)
  (let ((d (m+ m (transpose m))))
    (iter (for r from 0 below (matrix-rows d))
	  (never (not (plusp (mref d r r))))
	  (setq d (gaussian-pivot d r r)))))

(defun matrix-zero-p (m)
  (iter (for x in-vector m)
	(always (< (abs x) +epsilon+))))

(defun lexico-positive-p (v)
  "T if <v> is not the zero vector and the first non-zero element of
<v> is positive."
  (let ((val (iter (for x in-vector v)
		   (thereis (and (> (abs x) +epsilon+)
				 x)))))
    (and val (plusp val))))


(defun lexico-negative-p (v)
  "T if <v> is not the zero vector and the first nonzero element of
<v> is negative."
  (lexico-positive-p (mneg v)))

(defun lexico-positive-compare-p (a b)
  "T if a-b is lexico positive."
  (lexico-positive-p (m- a b)))

(defun lexico-negative-compare-p (a b)
  "T if a-b is lexico negative."
  (lexico-negative-p (m- a b)))

(defun lexico-minimum (vset)
  "Find the vector <vj> in the set of vectors <vset> with k vectors
that satisfies the property that all <vi> are lexographically greater
than <vj> for i = 1 to k."
  (first (sort vset #'lexico-negative-compare-p)))

(defun lexico-minimum-idx (vset)
    "Return the position of the vector <vj> in the set of vectors <vset> with k vectors
that satisfies the property that all <vi> are lexographically greater
than <vj> for i = 1 to k."
    (let* ((array-p (typep vset 'array))
	   (vj (if array-p
		   (get-column vset 0)
		   (first vset)))
	   (j 0))
      (if array-p
	  (iter (for k from 1 below (matrix-cols vset))
		(when (lexico-positive-compare-p vj (get-column vset k))
		  (setf vj (get-column vset k)
			j k))
		(finally (return j)))
	  (iter (for vk in (rest vset))
		(for k from 1)
		(when (lexico-positive-compare-p vj vk)
		  (setf vj vk
			j k))
		(finally (return j))))))

(defun positive-semidefinite-p (m) ;; Not working!
  (flet ((zeros-and-no-negs (m r-ignore)
	   (iter (for r from 0 below (matrix-rows m))
		 (unless (= r r-ignore)
		   (never (< 0 (mref m r r)))
		   (never (and (< (abs (mref m r r)) +epsilon+ )
			       (not (matrix-zero-p (get-column m r)))
			       (not (matrix-zero-p (get-row m r))))))))
	 (zero-diagonals (m)
	   (iter (for r from 0 below (matrix-rows m))
		 (when (< (mref m r r) +epsilon+)
		   (collect r)))))
    (let ((d (m+ m (transpose m))))
      (when (zeros-and-no-negs d 0)
	(let ((to-ignore (zero-diagonals d)))
	;; Check for 0 or negative diagonals
	(iter (for r from 0 below (matrix-rows d))
	      (always (zeros-and-no-negs d r))
	      (unless (< (abs (mref d r r)) +epsilon+)
		(setq d (gaussian-pivot d r r)))))))))

(defun %augment-matrix (m1 m2)
  ;(assert (= (matrix-rows m1) (matrix-rows m2)) (m1 m2) "Matrics ~S and ~S should have the same number of rows.")
  (let* ((km1 (matrix-cols m1))
	 (km2 (if (column-vector-p m2) 
		  1
		  (matrix-cols m2)))
	 (m (make-matrix (matrix-rows m1)
			 (+ km1 km2)
			 :element-type (element-type m1))))
    (iter (for i from 0 below (+ km1 km2))
	  (if (< i km1)
	      (set-column m i (get-column m1 i))
	      (set-column m i (if (column-vector-p m2)
				  m2
				  (get-column m2 (- i km1))))))
    m))

(defun augment-matrix (m1 m2 &rest m)
  (let ((am (%augment-matrix m1 m2)))
    (cond ((null m) am)
	  (t (apply #'augment-matrix am (first m) (rest m))))))

;;; ---- Inverting Matrices ----
(defun rank-3-inverse (m)
  "Invert a 3x3 using cofactors.  This method is faster than using a generic Gaussian elimination because of the loop overhead of such a method."
  (let* ((a11 (mref m 0 0)) (a12 (mref m 0 1)) (a13 (mref m 0 2))
	 (a21 (mref m 1 0)) (a22 (mref m 1 1)) (a23 (mref m 1 2))
	 (a31 (mref m 2 0)) (a32 (mref m 2 1)) (a33 (mref m 2 2))

	 (co11 (- (* a22 a33) (* a32 a23)))
	 (co12 (- (- (* a21 a33) (* a31 a23))))
	 (co13 (- (* a21 a32) (* a31 a22)))

	 (co21 (- (- (* a12 a33) (* a32 a13))))
	 (co22 (- (* a11 a33) (* a31 a13)))
	 (co23 (- (- (* a11 a32) (* a31 a12))))
	 
	 (co31 (- (* a12 a23) (* a22 a13)))
	 (co32 (- (- (* a11 a23) (* a21 a13))))
	 (co33 (- (* a11 a22) (* a21 a12)))
	 (det (/ 1 (- (+ (* a11 a22 a33)
			 (* a12 a23 a31)
			 (* a13 a21 a32))
		      (+ (* a13 a22 a31)
			 (* a12 a21 a33)
			 (* a11 a23 a32))))))
    (%matrix (list (list (* det co11) (* det co21) (* det co31))
		   (list (* det co12) (* det co22) (* det co32))
		   (list (* det co13) (* det co23) (* det co33))))))


(defun gauss-jordan-inverse (m)
  (let* ((a (augment-matrix m (identity-matrix (matrix-rows m))))
	 (ainv (make-matrix (matrix-rows m) (matrix-cols m))))
    (setq a (gaussian-forward-elimination a))
    (setq a (gaussian-back-substitution a))
    (iter (with k = (matrix-cols m))
	  (for i from 0 below k)
	  (set-column ainv i (get-column a (+ i k))))
    ainv))
    
(defun rank-2-inverse (m)
  (let* ((a (mref m 0 0))
	 (b (mref m 0 1))
	 (c (mref m 1 0))
	 (d (mref m 1 1))
	 (det (/ 1 (- (* a d) (* b c)))))
    (%matrix (list (list (* d det) (- (* b det)))
		   (list (- (* c det)) (* a det))))))
			 

(defmethod inverse ((m array))
  (let ((r (matrix-rows m))
	(c (matrix-cols m)))
    (cond ((= r c 2)
	   (rank-2-inverse m))
	  ((= r c 3)
	   (rank-3-inverse m))
	  (t (gauss-jordan-inverse m)))))

;;;; ------------------------------

(defun test-row-exchange ()
  (let ((m (%matrix '((1.0 -1.0 0.0)
		      (0.0 4.0 -1.0)
		      (2.0 0.0 1.0)))))
    (format t "Before: ~%")
    (print-matrix m)
    (format t "After: ~%")
    (print-matrix (gaussian-row-exchange m 0 1))))

(defun test-row-scale ()
  (let* ((m (%matrix '((1.0 -1.0 0.0)
		       (0.0 4.0 -1.0)
		       (2.0 0.0 1.0))))
	 (m1 (copy-matrix m))
	 (m2 (copy-matrix m)))
    (format t "Before: ~%")
    (print-matrix m)
    (format t "After: ~%")
    (print-matrix (gaussian-row-scale-slow m 0 10))
    (assert (matrix-equal (gaussian-row-scale-slow m1 0 10)
			  (gaussian-row-scale m2 0 10)))))

(defun test-row-scale-sum ()
  (let* ((m (%matrix '((1.0 -1.0 0.0)
		      (0.0 4.0 -1.0)
		       (2.0 0.0 1.0))))
	 (m1 (copy-matrix m))
	 (m2 (copy-matrix m)))
    (format t "Before: ~%")
    (print-matrix m)
    (format t "After: ~%")
    (print-matrix (gaussian-row-scale-sum-slow m 0 -2 2))
    (assert (matrix-equal (gaussian-row-scale-sum-slow m1 0 -2 2)
			  (gaussian-row-scale-sum m2 0 -2 2)))))


(defun test-forward-elim ()
  (let ((m (%matrix '((1 -1 0)
		      (0 4 -1)
		      (2 0 1)))))
    (format t "Before: ~%")
    (print-matrix m)
    (format t "After: ~%")
    (print-matrix (gaussian-forward-elimination m))))

(defun test-forward-elim2 ()
  (let ((m (%matrix '((1 1 1 6)
		      (1 2 2 11)
		      (2 3 -4 3)))))
    (print-matrix (gaussian-forward-elimination m))))

(defun test-forward-elim3 ()
  (let ((m (%matrix '((1 1 1 6)
		      (1 2 2 11)
		      (2 3 -4 3)))))
    (setf m (gaussian-forward-elimination m))
    (print-matrix (gaussian-back-substitution m))))


(defun test-non-square-mult ()
  (let ((m1 (%matrix '((1 0 -1)
		       (-1 2 0))))
	(m2 (%matrix '((-1 1 2)
		       (0 1 -3)
		       (-2 3 0)))))
    (format t "Is square? ~A~%" (square-matrix-p m1))
    (print-matrix (m* m1 m2))))

(defun test-lexico-positive ()
  (assert (lexico-positive-p (make-vector 43 5 -2 2 -34 5 -6)))
  (assert (lexico-positive-p (make-vector 0.0 0.0 0.0001 -1000.0)))
  (assert (lexico-negative-p (make-vector 0.0 -1.0 20.0 5000.0)))
  (assert (not (lexico-negative-p (make-vector 43.0 0.0 4.0 3.0 -2.0 2.0 -34.0 5.0 -6.0))))
  (assert (not (lexico-positive-p (make-vector -43.0 2.0 3.0 34.0 21.0 31.0 23.0 9.0 28.0))))
  (let ((set (list (make-vector -2 1 -20 -30)
		   (make-vector -2 0 -1 1)
		   (make-vector -2 0 -1 0)
		   (make-vector 0 -10 -40 -50))))
    (format t "~A~%" (lexico-minimum (copy-list set))) ;; Uses sort, which is destructive.
    (lexico-minimum-idx set)))

(defun test-qr ()
  (let ((a (%matrix '((1.0 1.0 2.0)
		      (0.0 0.0 1.0)
		      (1.0 0.0 0.0)))))
    (multiple-value-bind (q r) (qr-factorization a)
      (print-matrix q)
      (terpri)
      (print-matrix r))))

(defun test-tridiagonal-qr ()
  (let ((a (%matrix '((4.0 -3.0 0.0 0.0)
		      (-3.0 3.33 -1.67 0)
		      (0.0 -1.67 -1.32 0.91)
		      (0.0 0.0 0.91 1.99)))))
    (print-matrix (fast-qr-factorization a))
    (terpri)
    (print-matrix (orthogonalize-matrix a))))

